package com.cqs.leetcode.recursion;

import com.cqs.leetcode.ds.ArraysUtils;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

/**
 * 思路参考
 * https://leetcode-cn.com/problems/sum-of-distances-in-tree/solution/c-liang-ci-dfsde-dao-da-an-by-congwang357-2/
 * @author lixiaowen
 * @date 2020/10/6
 **/
public class SumDistancesTree834 {


    int N, sum = 0;
    int[] dist;
    int[] cnt;
    Map<Integer, Set<Integer>> graph = new HashMap<>();

    public int[] sumOfDistancesInTree(int N, int[][] edges) {
        this.N = N;
        //dp表示到点i的距离和
        dist = new int[N];
        cnt = new int[N];
        //假设节点0为根节点
        sum = 0;
        for (int i = 0; i < edges.length; i++) {
            Set<Integer> edge = graph.computeIfAbsent(edges[i][0], k -> new HashSet<>());
            edge.add(edges[i][1]);

            Set<Integer> edge2 = graph.computeIfAbsent(edges[i][1], k -> new HashSet<>());
            edge2.add(edges[i][0]);
        }
        rootDfs(0, -1, 1);
        dist[0] = sum;
        dfs(0, -1);
        return dist;
    }

    private void rootDfs(int node, int parent, int level) {
        //求root的孩子节点
        Set<Integer> edges = graph.get(node);
        if (edges != null) {
            for (Integer child : edges) {
                if (child == parent) continue;
                sum += level;
                rootDfs(child, node, level + 1);
                cnt[node] += cnt[child];
            }
        }
        cnt[node] += 1;
    }

    private void dfs(int node, int parent) {
        Set<Integer> edges = graph.get(node);
        if (edges != null) {
            for (Integer child : edges) {
                if (child == parent) continue;
                dist[child] = dist[node] + N - 2 * cnt[child];
                dfs(child, node);
            }
        }
    }


    public static void main(String[] args) {
        SumDistancesTree834 distances = new SumDistancesTree834();
        int N = 6;
//        int[][] edges = {{0, 1}, {0, 2}, {2, 3}, {2, 4}, {2, 5}};
        N = 3;
        int[][] edges = {{0, 2}, {1, 2}};
        int[] ints = distances.sumOfDistancesInTree(N, edges);
        System.out.println(ArraysUtils.formatInts(ints));
    }
}
